#include <iostream>

#include "utils.hpp"
#include "nets.hpp"

int main(int argc, char* argv[])
{
    std::cout << "Using Libtorch : " << TORCH_VERSION << std::endl;

    // 构造数据
    std::pair<std::vector<float>, std::vector<float>> train_data = createData();

    std::vector<float>& inputs = train_data.first;
    std::vector<float>& labels = train_data.second;

    // 容器转为张量
    auto inputs_tensor = torch::from_blob(inputs.data(), {int(inputs.size()), 1});
    auto labels_tensor = torch::from_blob(labels.data(), {int(inputs.size()), 1});

    auto dataSet =
        CustomDtaset(inputs_tensor, labels_tensor).map(torch::data::transforms::Stack<>());
    auto dataLoader = torch::data::make_data_loader<torch::data::samplers::RandomSampler>(
        std::move(dataSet), 16);

    // 创建神经网络
    auto net       = std::make_shared<Net>(1, 1);
    auto optimizer = std::make_shared<torch::optim::Adam>(net->parameters(), 0.001);

    // 训练并打印损失
    int n_epochs = 400;
    for (int epoch = 0; epoch <= n_epochs; epoch++)
    {
        float lossVal = 0.0;
        net->train();
        //
        for (auto& batch : *dataLoader)
        {
            auto pred = net->forward(batch.data);

            auto loss = torch::mse_loss(pred, batch.target);

            // 反向传播、计算梯度、更新参数
            optimizer->zero_grad();
            loss.backward();
            optimizer->step();

            lossVal += loss.item<float>();
        }

        // 打印训练信息
        if (epoch % (n_epochs / 5) == 0)
        {
            for (const auto& par : net->named_parameters())
            {
                std::cout << par.key() << "\t" << par.value().item<float>() << std::endl;
            }
            std::cout << "Epoch." << epoch << ", Loss: " << lossVal << std::endl << std::endl;
        }

        net->eval();
        auto out = net->forward(inputs_tensor);
        std::ofstream os("train_results_" + std::to_string(epoch) + ".csv");
        torch::Tensor utm1Temp = torch::zeros({1, 2});
        torch::Tensor xtm1Temp = torch::zeros({1, 2});
        torch::Tensor ztTemp   = torch::zeros({1, 2});
        for (int i = 0; i < int(inputs.size()); ++i)
        {
            os << inputs[i] << " " << labels[i] << " " << out[i][0].item<float>() << std::endl;
        }
        os.close();
    }

    return 0;
}
